/*
 * File: ExportNeo4J.java 
 *  
 * Created on 23.01.2011
 */
package de.graph.impex;

import java.util.Date;
import java.util.HashMap;
import java.util.Map;
import java.util.TreeMap;

import org.neo4j.graphdb.DynamicRelationshipType;
import org.neo4j.graphdb.RelationshipType;
import org.neo4j.graphdb.index.BatchInserterIndex;
import org.neo4j.index.impl.lucene.LuceneBatchInserterIndexProvider;
import org.neo4j.kernel.impl.batchinsert.BatchInserter;
import org.neo4j.kernel.impl.batchinsert.BatchInserterImpl;

import de.graph.common.Edge;
import de.graph.common.Graph;
import de.graph.common.Vertex;


public class Neo4JExporter extends AbstractImportExport implements Exporter {
    
    private final BatchInserter inserter;

    private final LuceneBatchInserterIndexProvider indexProvider;
    
    private final BatchInserterIndex vertexIndex;

    private final BatchInserterIndex edgeIndex;

    public Neo4JExporter(String dbName) {
        this(dbName, false);
    }
    
    public Neo4JExporter(String dbName, boolean forceCleanup) {
        super(dbName, forceCleanup);
        
        Map<String, String> storeProperties = new HashMap<String, String>();
        storeProperties.put("neostore.nodestore.db.mapped_memory", "10M"); // nr_of_nodes * 9 bytes        
        storeProperties.put("neostore.relationshipstore.db.mapped_memory", "40M"); // nr_of_relationships * 33 bytes.
        storeProperties.put("neostore.propertystore.db.mapped_memory", "100M");
        storeProperties.put("neostore.propertystore.db.strings.mapped_memory", "300M");
        storeProperties.put("neostore.propertystore.db.arrays.mapped_memory", "0M");

        this.inserter = new BatchInserterImpl(dbName, storeProperties);        
        this.indexProvider = new LuceneBatchInserterIndexProvider(inserter);

        // use exact analysing (no full text splitting)
        Map<String, String> indexProperties = new HashMap<String, String>();
        indexProperties.put("type", "exact");
        
        this.vertexIndex = indexProvider.nodeIndex("VertexIndex", indexProperties);
        this.edgeIndex = indexProvider.nodeIndex("EdgeIndex", indexProperties);
    }
    
    public void exportGraph(Graph graph) {
        exportGraph(graph, null);
    }
    
    public void exportGraph(Graph graph, Vertex root) {
        try {
            Map<Vertex, Long> vertexNeoIdMap = new HashMap<Vertex, Long>();
            Map<String, RelationshipType> relationNameMap = new HashMap<String, RelationshipType>();
                        
            // special handling for root node
            if (root != null) {
                vertexNeoIdMap.put(root, getReferenceNode(inserter));
            }
            
            // save graph data
            inserter.createNode(getProperties(graph));
            
            System.out.println("add vertices: "+ new Date());
            
            // create nodes except root
            for (Vertex vertex : graph.getVertices()) {
                if (vertex != root) {
                    Map<String, Object> properties = getProperties(vertex);
                    
                    long neoNodeId = inserter.createNode(properties);

                    vertexNeoIdMap.put(vertex, neoNodeId);

                    // index all properties
                    vertexIndex.add(neoNodeId, properties);
                }
            }
                        
            System.out.println("add edges: "+ new Date());
            
            // create edges
            for (Edge edge : graph.getEdges()) {
                long neoSourceId = vertexNeoIdMap.get(edge.getSourceVertex());
                long neoTargetId = vertexNeoIdMap.get(edge.getTargetVertex());

                String label = edge.getLabel();
                RelationshipType relationshipType =  relationNameMap.get(label);
                
                if (relationshipType == null) {
                    relationshipType = DynamicRelationshipType.withName(label);
                    
                    relationNameMap.put(label, relationshipType);
                }
                
                Map<String, Object> properties = getProperties(edge);
                long neoEdgeId = inserter.createRelationship(neoSourceId, neoTargetId, relationshipType, properties);

                // index all properties
                edgeIndex.add(neoEdgeId, properties);
            }
        }
        finally {
            shutdown();
        }
    }

    @Override
    protected void shutdown() {
        System.out.println("shutdown indexer: "+ new Date());
        
        indexProvider.shutdown();
        
        System.out.println("shutdown inserter: "+ new Date());
        
        inserter.shutdown();
    }

    private Map<String,Object> getProperties(Graph graph) {
        
        Map<String,Object> properties = new TreeMap<String, Object>();
        
        properties.put("IsDirected", graph.isDirected());
        properties.put("NodeType", Graph.class.getName());
        
        return properties;
    }

    private Map<String,Object> getProperties(Edge vertex) {
        
        Map<String,Object> properties = new HashMap<String, Object>();        
        for (String propertyKey : vertex.getPropertyKeys()) {
            Object propertyValue = vertex.getProperty(propertyKey);
            
            properties.put(propertyKey, propertyValue);
        }
        
        // remove internal keys
        properties.remove("source");
        properties.remove("target");
        
        return properties;
    }

    private Map<String,Object> getProperties(Vertex vertex) {
        
        Map<String,Object> properties = new HashMap<String, Object>();        
        for (String propertyKey : vertex.getPropertyKeys()) {
            Object propertyValue = vertex.getProperty(propertyKey);
            
            properties.put(propertyKey, propertyValue);
        }
        
        // remove internal keys
        properties.remove("inEdges");
        properties.remove("outEdges");
        
        return properties;
    }

    private long getReferenceNode(BatchInserter inserter) {

        long referenceNodeID = 0L;
        if (inserter.getReferenceNode() < referenceNodeID) {
            inserter.createNode(referenceNodeID, null);
        }
         
        return referenceNodeID;
    }
}
